import time
import uuid
import redis


class DistributedLock:
    def __init__(self, redis_client, lock_name, expire_time=30, retry_interval=0.1, retry_count=30):
        """
        分布式锁
        :param redis_client: Redis 客户端连接
        :param lock_name: 锁名称
        :param expire_time: 锁的过期时间(秒)
        :param retry_interval: 获取锁失败后的重试间隔(秒)
        :param retry_count: 最大重试次数
        """
        self.redis = redis_client
        self.lock_name = lock_name
        self.expire_time = expire_time
        self.retry_interval = retry_interval
        self.retry_count = retry_count
        self.identifier = str(uuid.uuid4())  # 唯一标识锁持有者

    def acquire(self):
        """获取分布式锁"""
        retry = 0
        while retry < self.retry_count:
            # 使用 SET 命令的 NX 和 PX 参数实现原子操作
            acquired = self.redis.set(
                self.lock_name,
                self.identifier,
                nx=True,
                px=int(self.expire_time * 1000)  # 转换为毫秒
            )
            if acquired:
                return True
            time.sleep(self.retry_interval)
            retry += 1
        return False

    def release(self):
        """释放分布式锁"""
        # 使用 Lua 脚本保证原子性：只有锁持有者才能释放锁
        unlock_script = """
        if redis.call("get", KEYS[1]) == ARGV[1] then
            return redis.call("del", KEYS[1])
        else
            return 0
        end
        """
        result = self.redis.eval(unlock_script, 1, self.lock_name, self.identifier)
        return result == 1

    def __enter__(self):
        if self.acquire():
            return self
        raise Exception("Failed to acquire lock")

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.release()


# 使用示例
if __name__ == "__main__":
    redis_client = redis.Redis(host='localhost', port=6379, db=5)

    # 使用上下文管理器
    try:
        with DistributedLock(redis_client, "my_resource_lock"):
            print("Lock acquired, do something...")
            time.sleep(10)  # 业务处理
    except Exception as e:
        print(e)

    # 或手动管理
    # lock = DistributedLock(redis_client, "my_resource_lock")
    # if lock.acquire():
    #     try:
    #         print("Lock acquired, do something...")
    #         time.sleep(10)
    #     finally:
    #         lock.release()
    # else:
    #     print("Failed to acquire lock")
